package com.atguigu.java.ai.langchain4j;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.onnx.HuggingFaceTokenizer;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.apache.ibatis.javassist.ClassPath;
import org.junit.jupiter.api.Test;
import org.springframework.boot.test.context.SpringBootTest;

import java.nio.file.FileSystems;
import java.nio.file.PathMatcher;
import java.util.List;

@SpringBootTest
public class RAGTest {

    @Test
    public void testReadDocument(){

        //使用FileSystemDocumentLoader读取指定目录下的知识库文档
        //并使用默认的文档解析器TextDocumentParser对文档进行解析
        Document document = ClassPathDocumentLoader.loadDocument("knowledge\\测试.txt");
        System.out.println(document.text());

        // 加载单个文档
        Document document1 = ClassPathDocumentLoader.loadDocument("knowledge/医院信息.txt", new TextDocumentParser());

        // 从一个目录中加载所有文档
        List<Document> documents2 = ClassPathDocumentLoader.loadDocuments("knowledge", new TextDocumentParser());

        // 从一个目录中加载所有的.txt文档
        PathMatcher pathMatcher = FileSystems.getDefault().getPathMatcher("glob:*.txt");
        List<Document> documents = ClassPathDocumentLoader.loadDocuments("knowledge", pathMatcher, new TextDocumentParser());

        // 从一个目录及其子目录中加载所有文档
        List<Document> documents1 = ClassPathDocumentLoader.loadDocumentsRecursively("knowledge", new TextDocumentParser());

    }

    @Test
    public void  testParsePDF(){
        Document document =
                FileSystemDocumentLoader.loadDocument("D:\\gmallparent\\java-ai-langchain4j\\src\\main\\resources\\knowledge\\医院信息.pdf",new ApachePdfBoxDocumentParser());
        System.out.println(document);
    }


    /**
     * 加载一个测试文档
     * 默认的递归分割器，用文档分割器，将文档分成若干个段落
     * 默认的向量大模型
     * 把段落存储到向量数据库里面
     */
    @Test
    public void testReadDocumentAndStore(){

        Document document = ClassPathDocumentLoader.loadDocument("knowledge\\人工智能.md",new TextDocumentParser());

        //创一个向量数据库
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        //使用默认的文档分割器，将文档分割，转成向量，存储到向量数据库里面
        EmbeddingStoreIngestor.ingest(document,embeddingStore);

        System.out.println(embeddingStore);

    }


    @Test
    public void testDocumentSplitter(){
        //加载一个测试文档
        Document document = ClassPathDocumentLoader.loadDocument("knowledge\\人工智能.md",new TextDocumentParser());

        //创建文档分割器   段落文档分割器    每个短路300token   每个段落有30个重复的token
        DocumentByParagraphSplitter documentByParagraphSplitter = new DocumentByParagraphSplitter(
                300,
                30,
                new HuggingFaceTokenizer());

        //基于内存的文档向量数据库
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        //向量数据库中保存文档
        EmbeddingStoreIngestor.builder()
                .embeddingStore(embeddingStore)
                .documentSplitter(documentByParagraphSplitter)
                .build()
                .ingest(document);

    }


    @Test
    public void testTokenCount(){
        String text = "这是一个示例文本，用于测试token长度计算";
        UserMessage userMessage = UserMessage.userMessage(text);
        HuggingFaceTokenizer tokenizer = new HuggingFaceTokenizer();
        int count = tokenizer.estimateTokenCountInMessage(userMessage);
        System.out.println("Token count" + count);
    }



}
